/*
Copyright (C) 2011 by the PSVN Research Group, University of Alberta
*/

#include <stdlib.h>
#include <stdio.h>
#include <inttypes.h>
#include <assert.h>
#include <sys/time.h>


/* Always search to a fixed depth, ignoring the goal test.
   This is almost certainly not what you want, unless you
   are doing some sort of performance tests */
//#define DEPTH_LIMITED_SEARCH 8
/* If we're doing depth limited search, also do the goal computation
   and heuristic lookup even though we ignore the results */
//#define DO_GOAL_HEUR_COMPUTATION

/* if undefined, don't use the feature */
//#define PARENT_PRUNING
//#define MOVE_PRUNING
#define TTABLE_MB 1024

//#define PRINT_SOLUTION


#if defined( TTABLE_MB ) && defined( MOVE_PRUNING )
#define HISTORY_TT
#endif
#if defined( TTABLE_MB ) && defined( DEPTH_LIMITED_SEARCH )
#error TTABLE_MB and DEPTH_LIMITED_SEARCH should not both be on
#undef TTABLE_MB
#endif

static int min(int a, int b)
{
    return (a < b) ? a : b;
}

typedef struct {
  state_t state;
#ifdef HISTORY_TT
  int history;
#endif
  int h_value;		/* if state has been searched before, we now
			   know a (potentially) better estimate of
			   the true distance to goal.  -1 if entry is unused */
} tt_entry_t;

typedef struct {
  int32_t size;		/* number of tt entries */
  tt_entry_t *entries;
} ttable_t;

uint64_t nodes_visited;
#ifdef DO_GOAL_HEUR_COMPUTATION
uint64_t useless_value = 0;
#endif


void clear_ttable( ttable_t *ttable )
{
  int32_t i;

  for( i = 0; i < ttable->size; ++i ) {

    ttable->entries[ i ].h_value = -1;
  }
}

ttable_t *new_ttable( const int32_t sizeMB )
{
  ttable_t *ttable;

  ttable = malloc( sizeof( *ttable ) );
  assert( ttable != NULL );

  ttable->size = (int64_t)sizeMB * 1048576 / sizeof( tt_entry_t );
  if( ttable->size < 1 ) {
    ttable->size = 1;
  }
  ttable->entries = malloc( sizeof( ttable->entries[ 0 ] ) * ttable->size );
  assert( ttable->entries != NULL );

  clear_ttable( ttable );

  return ttable;
}

/* update the transposition table with a new lower bound estimate */
void update_ttable( ttable_t *ttable, const state_t *state,
#ifdef HISTORY_TT
		    const int history,
#endif
		    const int h_value )
{
  int32_t idx;

#ifdef HISTORY_TT
  idx = hash_state_history( state, history ) % ttable->size;
#else
  idx = hash_state( state ) % ttable->size;
#endif
  if( compare_states( state, &ttable->entries[ idx ].state ) == 0
#ifdef HISTORY_TT
      && ttable->entries[ idx ].history == history
#endif
      ) {
    /* state is in table */

    if( h_value > ttable->entries[ idx ].h_value ) {
      /* h_value is a larger lower bound */

      ttable->entries[ idx ].h_value = h_value;
    }
  } else {
    /* state wasn't in table - overwrite */

    copy_state( &ttable->entries[ idx ].state, state );
#ifdef HISTORY_TT
    ttable->entries[ idx ].history = history;
#endif
    ttable->entries[ idx ].h_value = h_value;
  }
}

/* get the lower bound estimate for a state
   returns -1 if the state is not in the table */
int check_ttable( const ttable_t *ttable, const state_t *state
#ifdef HISTORY_TT
		  , const int history
#endif
		  )
{
  int32_t idx;

#ifdef HISTORY_TT
  idx = hash_state_history( state, history ) % ttable->size;
#else
  idx = hash_state( state ) % ttable->size;
#endif
  if( compare_states( state, &ttable->entries[ idx ].state ) != 0
#ifdef HISTORY_TT
      || ttable->entries[ idx ].history != history
#endif
      ) {
    /* state not in table */

    return -1;
  }

  return ttable->entries[ idx ].h_value;
}

typedef struct {
    abstraction_t* abst;
    state_map_t* map;
} abstraction_data_t;

abstraction_data_t* read_abstraction_data(char* filename)
{
    abstraction_data_t* data;
    data = (abstraction_data_t*) malloc (sizeof(abstraction_data_t));
    data->abst = read_abstraction_from_file( filename );
    if( data->abst == NULL ) {
        fprintf( stderr, "could not read abstraction file %s\n", filename );
        return NULL;
    }
    /* get map filename by changing extension from ".abst" to ".table" */
    char map_filename[1024];
    strcpy(map_filename, filename);
    strcpy(strstr(map_filename, ".abst"), ".table");
    
    /* read the map */
    FILE* file = fopen( map_filename, "r" );
    if( file == NULL ) {
        fprintf( stderr, "could not open abstraction table %s\n", map_filename);
        destroy_abstraction( data->abst );
        free( data );
        return NULL;
    }
    data->map = read_state_map( file );
    fclose( file );
    if( data->map == NULL ) {
        fprintf( stderr, "could not read abstraction table %s\n", map_filename);
        destroy_abstraction( data->abst );
        free( data );
        return NULL;
    }
    return data;
}

void destroy_abstraction_data(abstraction_data_t* data)
{
    destroy_abstraction(data->abst);
    destroy_state_map(data->map);
    free(data);
}

int abstraction_data_lookup(const abstraction_data_t* abst, 
                            const state_t* state)
{
    state_t abst_state;
    abstract_state( abst->abst, state, &abst_state );
    int *h;
    h = state_map_get( abst->map, &abst_state );
    if (h == NULL) {
        return 0;
    }
    return *h;
}

int dfs( const state_t *state,
#ifdef PARENT_PRUNING
	 const state_t *parent_state,
#endif
#ifdef MOVE_PRUNING
	 const int history,
#endif
#ifdef TTABLE_MB
	 ttable_t *ttable,
#endif
	 const int max_cost,
         int *min_need )
{

    int rule_used;
#ifdef MOVE_PRUNING
    int c_history;
#endif
    func_ptr iter;
    state_t child;

    ++nodes_visited;

#ifdef DEPTH_LIMITED_SEARCH
#ifdef DO_GOAL_HEUR_COMPUTATION
    if( is_goal( state, game ) ) { ++useless_value; }
#endif
    if( max_cost == 0 ) {

        return 0;
    }
#else

#ifdef TTABLE_MB
    if( (*min_need = check_ttable( ttable, state
#ifdef HISTORY_TT
                                  , history
#endif
                                   )) > max_cost ) {
        return 0;
    } 
#endif
    if ( max_cost == 0 ) {
        if( is_goal( state ) ) {
#ifdef PRINT_SOLUTION
            print_state( stdout, state );
            printf( "\n" );
#endif
            return 1;
        }
        *min_need = COST_OF_CHEAPEST_FWD_RULE;
        return 0;
    }
#endif

    *min_need = INT_MAX;
    init_fwd_iter( iter );
    while( ( rule_used = next_fwd_iter( iter, state ) ) >= 0 ) 
    {
#ifdef MOVE_PRUNING
        if( !fwd_rule_valid_for_history( history, rule_used ) )
            continue;
        c_history = next_fwd_history( history, rule_used );
#endif
        apply_fwd_rule( rule_used, state, &child );
#ifdef PARENT_PRUNING
        if( compare_states( &child, parent_state ) == 0 )
            continue;
#endif

#ifdef DEPTH_LIMITED_SEARCH
        const int move_cost = 1;
#else
        const int move_cost = fwd_rule_costs[ rule_used ];
#endif
        if ( move_cost > max_cost )
            *min_need = min( *min_need, move_cost );
        else {
            int child_need;
            if( dfs( &child,
#ifdef PARENT_PRUNING
                     state,
#endif
#ifdef MOVE_PRUNING
                     c_history,
#endif
#ifdef TTABLE_MB
                     ttable,
#endif
                     max_cost - move_cost, &child_need ) ) 
            {
#ifdef PRINT_SOLUTION
                printf( "%s\n", fwd_rule_names[ rule_used ] );
                print_state( stdout, state );
                printf( "\n" );
#endif
                return 1;
            }
            *min_need = min( *min_need, move_cost + child_need );
        }
    }
    assert( *min_need > max_cost );
#ifdef TTABLE_MB
#ifdef HISTORY_TT
    update_ttable( ttable, state, history, *min_need );
#else
    update_ttable( ttable, state, *min_need );
#endif
#endif
    return 0;
}

int dfs_heur( const abstraction_data_t *abst,
	      const state_t *state,
#ifdef PARENT_PRUNING
	      const state_t *parent_state,
#endif
#ifdef MOVE_PRUNING
	      const int history,
#endif
#ifdef TTABLE_MB
	      ttable_t *ttable,
#endif
	      const int max_cost, int *min_need )
{
    int rule_used, h;
#ifdef TTABLE_MB
    int h2;
#endif
#ifdef MOVE_PRUNING
    int c_history;
#endif
    func_ptr iter;
    state_t child;

    ++nodes_visited;

#ifdef DEPTH_LIMITED_SEARCH
#ifdef DO_GOAL_HEUR_COMPUTATION
    if( is_goal( state ) ) { ++useless_value; }
    abstraction_data_lookup( abst, state );
#endif
    if( max_cost == 0 )
        return 0;
#else

    h = abstraction_data_lookup( abst, state );
#ifdef TTABLE_MB
    h2 = check_ttable( ttable, state
#ifdef HISTORY_TT
                       , history
#endif
                       );
    if( h2 > h ) {
        h = h2;
    }
#endif

    if( h > max_cost ) {
        *min_need = h;
        return 0;
    } else if ( max_cost == 0 ) {
        if( is_goal( state ) ) {
#ifdef PRINT_SOLUTION
            print_state( stdout, state );
            printf( "\n" );
#endif
            return 1;
        } else {
            *min_need = COST_OF_CHEAPEST_FWD_RULE;
            return 0;
        }
    }
#endif

    *min_need = INT_MAX;
    init_fwd_iter( iter );
    while( ( rule_used = next_fwd_iter( iter, state ) ) >= 0 ) {
#ifdef MOVE_PRUNING
        if( !fwd_rule_valid_for_history( history, rule_used ) )
            continue;
        c_history = next_fwd_history( history, rule_used );
#endif
        apply_fwd_rule( rule_used, state, &child );
#ifdef PARENT_PRUNING
        if( compare_states( &child, parent_state ) == 0 )
            continue;
#endif

#ifdef DEPTH_LIMITED_SEARCH
        const int move_cost = 1;
#else
        const int move_cost = fwd_rule_costs[ rule_used ];
#endif
        
        if (move_cost > max_cost)
            *min_need = min( *min_need, move_cost );
        else {
            int child_need;
            if( dfs_heur( abst, &child,
#ifdef PARENT_PRUNING
                          state,
#endif
#ifdef MOVE_PRUNING
                          c_history,
#endif
#ifdef TTABLE_MB
                          ttable,
#endif
                          max_cost - move_cost, &child_need ) ) 
            {
#ifdef PRINT_SOLUTION
                printf( "%s\n", fwd_rule_names[ rule_used ] );
                print_state( stdout, state );
                printf( "\n" );
#endif
                return 1;
            }
            *min_need = min( *min_need, move_cost + child_need );
        }
    }
    assert( *min_need > max_cost );
#ifdef TTABLE_MB
#ifdef HISTORY_TT
    update_ttable( ttable, state, history, *min_need );
#else
    update_ttable( ttable, state, *min_need );
#endif
#endif
  return 0;
}

int idastar( const abstraction_data_t *abst, const state_t *state,
#ifdef TTABLE_MB
	     ttable_t *ttable,
#endif
	     int64_t *total_nodes )
{
    int max_cost, done;

    *total_nodes = 0;
    if( abst != NULL ) {
        max_cost = abstraction_data_lookup( abst, state );
    } else {
        max_cost = 0;
    }
    while (1) {
        int need;
        nodes_visited = 0;
        if( abst != NULL ) {
            done = dfs_heur( abst, state,
#ifdef PARENT_PRUNING
                             state,
#endif
#ifdef MOVE_PRUNING
                             init_history,
#endif
#ifdef TTABLE_MB
                             ttable,
#endif
                             max_cost, &need );
        } else {
            done = dfs( state,
#ifdef PARENT_PRUNING
                        state,
#endif
#ifdef MOVE_PRUNING
                        init_history,
#endif
#ifdef TTABLE_MB
                        ttable,
#endif
                        max_cost, &need );
        }
        printf( "search cost %d: %ld\n", max_cost, nodes_visited );
        *total_nodes += nodes_visited;
        if( done ) {
            break;
        }
        assert( need > max_cost );
        max_cost = need;
    }
    return max_cost;
}

int main( int argc, char **argv )
{
    state_t state;
    int64_t visited, total_visited;
    int trials, d, total_d, need;
    abstraction_data_t* abst;
#ifdef TTABLE_MB
    ttable_t *ttable;
#endif
    char line[ 4096 ];
    struct timeval start, end;
    
    /* read the abstraction and value map */
    if( argc > 1 ) {
        abst = read_abstraction_data( argv[1] );
        if (abst == NULL) {
            return EXIT_FAILURE;
        }
    } else {
        
        abst = NULL;
    }

#ifdef TTABLE_MB
    ttable = new_ttable( TTABLE_MB );
#endif

    //srandom( 7 );
    total_d = 0;
    total_visited = 0;
    gettimeofday( &start, NULL );
    for( trials = 0;
         fgets( line, 4096, stdin ) != NULL
             && read_state( line, &state ) > 0;
         ++trials ) {
        
        printf( "problem %d: ", trials + 1 );
        print_state( stdout, &state );
        printf( "\n" );

#ifdef DEPTH_LIMITED_SEARCH
        d = DEPTH_LIMITED_SEARCH;
        nodes_visited = 0;
        if( abst != NULL ) {

            dfs_heur( abst, &state,
#ifdef PARENT_PRUNING
                      &state,
#endif
#ifdef MOVE_PRUNING
                      init_history,
#endif
#ifdef TTABLE_MB
                      ttable,
#endif
                      d, &need );
        } else {

            dfs( &state,
#ifdef PARENT_PRUNING
                 &state,
#endif
#ifdef MOVE_PRUNING
                 init_history,
#endif
#ifdef TTABLE_MB
                 ttable,
#endif
                 d, &need );
        }
        visited = nodes_visited;

#else /* DEPTH_LIMITED_SEARCH */
        need = 0;
        d = idastar( abst, &state,
#ifdef TTABLE_MB
                     ttable,
#endif
                     &visited );
#endif
        printf( "depth %d nodes %lu\n", d, visited );
        total_d += d;
        total_visited += visited;
    }
    gettimeofday( &end, NULL );
    
    end.tv_sec -= start.tv_sec;
    end.tv_usec -= start.tv_usec;
    if( end.tv_usec < 0 ) {
        end.tv_usec += 1000000;
        --end.tv_sec;
    }
    
    printf( "total: depth %d nodes %lu %zd.%06zd seconds\n",
            total_d, total_visited, end.tv_sec, end.tv_usec );
    
    if( abst != NULL ) {
        destroy_abstraction_data( abst );
    }
    
    return EXIT_SUCCESS;
}
